from langchain.chat_models import ChatOpenAI
from PyCmpltrtok.common import sep

chat_model = ChatOpenAI(
    streaming=True,
    verbose=True,
    callbacks=[],
    openai_api_key="token1",
    openai_api_base=f"http://127.0.0.1:6001/v1",
    model_name="chatglm2-6b",
    temperature=1.0,
    openai_proxy=None,
    top_p=1.0,
    max_tokens=2048,
)

if '__main__' == __name__:
    
    from langchain.schema import HumanMessage, SystemMessage, AIMessage

    messages = [
        SystemMessage(content='我们来玩成语接龙。我先来。'),
        AIMessage(content='为所欲为'),
        HumanMessage(content='围魏救赵'),
    ]
    
    xresult = chat_model.invoke(messages)
    sep()
    print(xresult)
    sep()
    print(xresult.content)
    sep()
